from __future__ import print_function
import torch.utils.data as data
from PIL import Image
import os
import os.path
import errno
import numpy as np
import torch
import codecs
# import skimage.data
import skimage.transform
import skimage.io


def load_data(data_dir, train=True):
    """Loads a data set and returns two lists:
    images: a list of Numpy arrays, each representing an image.
    labels: a list of numbers that represent the images labels.
    """
    # 初始化变量，确保在train=False时也能访问
    images = []
    labels = []

    if train:
        # 原始训练数据处理代码
        directories = [d for d in os.listdir(data_dir)
                       if os.path.isdir(os.path.join(data_dir, d))]
        for d in directories:
            label_dir = os.path.join(data_dir, d)
            file_names = [os.path.join(label_dir, f)
                          for f in os.listdir(label_dir) if f.endswith(".png")]
            for f in file_names:
                images.append(skimage.io.imread(f))
                labels.append(int(d))
    else:
        # 添加测试数据处理代码
        # 假设测试数据目录结构相同
        # 🚨 假设测试集没有标签，仅做分类预测（或统一标签为 0 占位）
        file_names = [os.path.join(data_dir, f)
                      for f in os.listdir(data_dir) if f.endswith(".png")]
        for f in file_names:
            images.append(skimage.io.imread(f))
            labels.append(0)  # dummy label，仅占位

    # 后续处理保持不变
    images28 = [skimage.transform.resize(image, (28, 28)) for image in images]
    labels_a = np.asarray(labels)
    images_a = np.asarray(images28)
    return images_a, labels_a


class Traffic(data.Dataset):
    '''
    Traffic Dataset.
    '''

    def __init__(self, root, train=True, transform=None, target_transform=None, download=False):
        self.root = os.path.expanduser(root)
        self.transform = transform
        self.target_transform = target_transform
        self.train = train  # training set or test set

        if self.train:
            self.train_data, self.train_labels = load_data(self.root, train=True)
        else:
            self.test_data, self.test_labels = load_data(self.root, train=False)

    def __getitem__(self, index):
        """
        Args:
            index (int): Index

        Returns:
            tuple: (image, target) where target is index of the target class.
        """
        if self.train:
            img, target = self.train_data[index], self.train_labels[index]
        else:
            img, target = self.test_data[index], self.test_labels[index]

        # doing this so that it is consistent with all other datasets
        img = img.astype(np.float32)
        target = torch.LongTensor([target])[0]

        if self.transform is not None:
            img = self.transform(img)

        if self.target_transform is not None:
            target = self.target_transform(target)

        return img, target

    def __len__(self):
        if self.train:
            return len(self.train_data)
        else:
            return len(self.test_data)